import pandas as pd
import numpy as np
import requests
from bs4 import BeautifulSoup
import time
import re
import urllib.parse
import os
import random
import traceback
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

def create_session_with_retry():
    """
    Create a session with retry capabilities
    """
    session = requests.Session()
    
    # Configure retry strategy
    retries = Retry(
        total=5,
        backoff_factor=1,
        status_forcelist=[429, 500, 502, 503, 504],
        allowed_methods=["GET", "POST"]
    )
    
    # Mount the adapter to the session
    adapter = HTTPAdapter(max_retries=retries)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    
    return session

def get_browser_headers():
    """
    Return realistic browser headers
    """
    return {
        'User-Agent': 'Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1',
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
        'Accept-Language': 'zh-CN,zh;q=0.9',
        'Connection': 'keep-alive'
    }

def extract_chemical_info(html_content):
    """
    Extract chemical information from the boxList div
    """
    # Initialize properties dictionary
    properties = {
        "中文名称": "",
        "英文名称": "",
        "分子式": "",
        "CAS号": "",
        "活性": "",
        "药理作用": "",
        "CB号": ""
    }
    
    # Parse HTML
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Find the boxList div
    box_list = soup.find('div', id='boxList')
    if not box_list:
        print("boxList div not found")
        return properties
    
    # Extract list items
    list_items = box_list.find_all('li')
    
    for item in list_items:
        # Extract raw text and remove extra whitespace
        item_text = item.get_text(strip=True)
        
        # Extract CB number (CBNumber)
        cb_link = item.select_one('a[href*="ProductChemicalPropertiesCB"]')
        if cb_link:
            cb_match = re.search(r'ProductChemicalPropertiesCB(\d+)', cb_link['href'])
            if cb_match:
                properties["CB号"] = f"CB{cb_match.group(1)}"
        
        # Extract English name
        if '英文名称' in item_text:
            # The English name is the text after '英文名称'
            english_name = item_text.replace('英文名称', '', 1).strip()
            properties["英文名称"] = english_name
        
        # Extract Chinese name
        if '中文名称' in item_text:
            # The Chinese name is the text after '中文名称'
            chinese_name = item_text.replace('中文名称', '', 1).strip()
            properties["中文名称"] = chinese_name
        
        # Extract molecular formula (MF)
        if 'MF' in item_text:
            # The formula is the text after 'MF'
            formula = item_text.replace('MF', '', 1).strip()
            properties["分子式"] = formula
            
        # Extract CAS number
        if item_text.startswith('CAS'):
            # The CAS number is the text after 'CAS'
            cas = item_text.replace('CAS', '', 1).strip()
            properties["CAS号"] = cas
    
    return properties

def get_additional_chemical_details(session, cb_number):
    """
    Get additional chemical details from the CB number URL
    """
    # Default empty additional properties
    additional_properties = {
        "活性": "",
        "药理作用": ""
    }
    
    if not cb_number:
        return additional_properties
    
    # Construct the full URL
    details_url = f"https://m.chemicalbook.com/ProductChemicalPropertiesCB{cb_number}.htm"
    
    try:
        # Add random delay to mimic human behavior
        time.sleep(random.uniform(2, 5))
        
        # Send request
        response = session.get(details_url, headers=get_browser_headers(), timeout=30)
        
        # Check response
        if response.status_code != 200:
            print(f"Failed to fetch details for {cb_number}: HTTP {response.status_code}")
            return additional_properties
        
        # Parse HTML
        soup = BeautifulSoup(response.text, 'html.parser')
        
        # Extract biological activity
        biological_activity_div = soup.find('div', class_='gsbt', string=re.compile('生物活性'))
        if biological_activity_div and biological_activity_div.next_sibling:
            biological_activity = biological_activity_div.next_sibling.strip()
            additional_properties["活性"] = biological_activity
        
        # Extract pharmacological effects
        pharmacological_effects_div = soup.find('div', class_='gsbt', string=re.compile('用途'))
        if pharmacological_effects_div and pharmacological_effects_div.next_sibling:
            pharmacological_effects = pharmacological_effects_div.next_sibling.strip()
            additional_properties["药理作用"] = pharmacological_effects
        
        return additional_properties
    
    except Exception as e:
        print(f"Error fetching details for {cb_number}: {e}")
        return additional_properties

def search_chemicalbook_by_cas(session, cas_number):
    """
    Search ChemicalBook using CAS number
    """
    # Initialize an empty properties dictionary
    empty_properties = {
        "中文名称": "",
        "英文名称": "",
        "分子式": "",
        "CAS号": "",
        "活性": "",
        "药理作用": "",
        "CB号": ""
    }
    
    if not cas_number or pd.isna(cas_number) or cas_number == "":
        return empty_properties
    
    # Clean the CAS number
    cas_number = str(cas_number).strip()
    
    # Encode the CAS number for the URL
    encoded_cas = urllib.parse.quote(cas_number)
    search_url = f"https://m.chemicalbook.com/Search.aspx?keyword={encoded_cas}"
    
    try:
        # Add random delay to mimic human behavior
        time.sleep(random.uniform(2, 5))
        
        # Send request
        response = session.get(search_url, headers=get_browser_headers(), timeout=30)
        
        # Check response
        if response.status_code != 200:
            print(f"Failed to search for CAS '{cas_number}': HTTP {response.status_code}")
            return empty_properties
        
        # Extract chemical information
        properties = extract_chemical_info(response.text)
        
        # Verify the CAS number matches
        if properties["CAS号"] and properties["CAS号"].strip() != cas_number:
            print(f"CAS number mismatch: searched for '{cas_number}', but found '{properties['CAS号']}'")
            return empty_properties
            
        return properties
    
    except Exception as e:
        print(f"Error searching for CAS '{cas_number}': {e}")
        return empty_properties

def process_excel_file(input_file, output_file):
    """
    Process the input Excel file, search ChemicalBook for each CAS number
    Skip rows without a CAS number or if no match is found
    Add a "随便" row at the beginning and remove it before saving
    """
    # Read the input Excel file
    df = pd.read_excel(input_file)
    
    # Check if the required column exists
    if "CAS号" not in df.columns:
        print("Error: 'CAS号' column not found in the Excel file.")
        print(f"Available columns: {', '.join(df.columns)}")
        return
    
    # Add the "随便" row at the beginning
    print("Adding '随便' row at the beginning for API compatibility")
    new_row = pd.DataFrame({
        col: ['随便' if col == '中文名' else 
              '9999-9999-9999' if col == 'CAS号' else 
              None] 
        for col in df.columns
    }, index=[0])
    df = pd.concat([new_row, df]).reset_index(drop=True)
    
    # Create a session for all requests
    session = create_session_with_retry()
    
    # Iterate through each row in the input DataFrame
    for index, row in df.iterrows():
        cas_number = row.get("CAS号", "")
        
        # Skip rows without a CAS number
        if pd.isna(cas_number) or cas_number == "":
            print(f"Skipping row {index+1}: No CAS number found")
            continue
        
        print(f"Processing {index+1}/{len(df)}: CAS={cas_number}")
        
        # Search ChemicalBook using the CAS number
        properties = search_chemicalbook_by_cas(session, cas_number)
        
        # Skip if no match found (empty properties)
        if not properties["CAS号"]:
            print(f"Skipping row {index+1}: No match found for CAS={cas_number}")
            continue
        
        # If CB号 exists, fetch additional details
        if properties["CB号"]:
            cb_number = properties["CB号"].replace("CB", "")
            additional_details = get_additional_chemical_details(session, cb_number)
            
            # Merge the additional details with existing properties
            properties.update(additional_details)
        
        # Update Chinese name if it's missing in the original data
        if properties["中文名称"] and "中文名" in df.columns:
            current_cn_name = row.get("中文名", "")
            if pd.isna(current_cn_name) or current_cn_name == "":
                df.at[index, "中文名"] = properties["中文名称"]
                print(f"  Added Chinese name: {properties['中文名称']}")
        
        # Update activity, pharmacological effect, and CB number
        if "活性" in df.columns and properties["活性"]:
            df.at[index, "活性"] = properties["活性"]
            
        if "药理作用" in df.columns and properties["药理作用"]:
            df.at[index, "药理作用"] = properties["药理作用"]
            
        if "CB号" in df.columns and properties["CB号"]:
            df.at[index, "CB号"] = properties["CB号"]
    
    # Remove the "随便" row before saving
    if len(df) > 0 and df.iloc[0]['中文名'] == '随便':
        print("Removing the '随便' row before saving")
        df = df.iloc[1:].reset_index(drop=True)
    
    # Save the results to the output Excel file
    df.to_excel(output_file, index=False)
    print(f"Results saved to {output_file}")

def main():
    input_file = "./API/mjz0320_matched_PC.xlsx"
    output_file = "./API/mjz0320_ChemicalBook_结果.xlsx"
    
    # Check if the input file exists
    if not os.path.exists(input_file):
        print(f"Input file '{input_file}' not found.")
        return
    else:
        print(f"Found input file: {input_file}")
    
    # Process the file
    process_excel_file(input_file, output_file)

if __name__ == "__main__":
    main()